import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from enum import Enum

# 飞行模式枚举
class FlightMode(Enum):
    TAKEOFF = 0
    WAYPOINT = 1
    RETURN_HOME = 2
    LANDING = 3

# 测试模式枚举
class TestMode(Enum):
    TAKEOFF_TEST = 0
    LANDING_TEST = 1
    WAYPOINT_TEST = 2
    RETURN_HOME_TEST = 3
    COMPLETE_TEST = 4

# 四旋翼参数
MASS = 1.0         # kg
GRAVITY = 9.81     # m/s^2
DT = 0.01          # 时间步长
SIM_TIME = 30      # 仿真时长（秒）
INERTIA = np.array([0.1, 0.1, 0.1])  # 转动惯量

# 飞行参数
TAKEOFF_HEIGHT = 12.0  # 起飞高度
LANDING_SPEED = 0.5   # 降落速度
WAYPOINT_RADIUS = 0.1 # 到达航路点的判定半径
HOME_POSITION = np.array([0.0, 0.0, 0.0])  # 返航点位置

# 航路点列表
WAYPOINTS = [
    np.array([2.0, 2.0, 2.0]),
    np.array([2.0, -2.0, 2.0]),
    np.array([-2.0, -2.0, 2.0]),
    np.array([-2.0, 2.0, 2.0])
]

# 四旋翼几何参数
ARM_LENGTH = 0.2   # 旋翼臂长
ROTOR_RADIUS = 0.05  # 旋翼半径
AXIS_LENGTH = 0.3  # 姿态指示器长度

# 测试参数
TEST_COMPLETION_TIME = 5.0  # 每个测试模式完成后等待时间
test_mode = TestMode.WAYPOINT_TEST  # 当前测试模式
test_completion_time = 0  # 测试完成计时器

def create_quadcopter_vertices(pos, att):
    """创建四旋翼的顶点"""
    # 创建基础顶点（未旋转）
    vertices = np.array([
        [ARM_LENGTH, 0, 0],    # 右前
        [0, ARM_LENGTH, 0],    # 前
        [-ARM_LENGTH, 0, 0],   # 左前
        [0, -ARM_LENGTH, 0],   # 后
    ])
    
    # 旋转矩阵
    roll, pitch, yaw = att
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(roll), -np.sin(roll)],
                   [0, np.sin(roll), np.cos(roll)]])
    
    Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                   [0, 1, 0],
                   [-np.sin(pitch), 0, np.cos(pitch)]])
    
    Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                   [np.sin(yaw), np.cos(yaw), 0],
                   [0, 0, 1]])
    
    R = Rz @ Ry @ Rx
    
    # 应用旋转
    vertices = vertices @ R.T
    
    # 添加位置偏移
    vertices += pos
    
    return vertices

def create_attitude_indicators(pos, att):
    """创建姿态指示器"""
    # 创建基础坐标轴
    axes = np.array([
        [AXIS_LENGTH, 0, 0],    # X轴
        [0, AXIS_LENGTH, 0],    # Y轴
        [0, 0, AXIS_LENGTH],    # Z轴
    ])
    
    # 旋转矩阵
    roll, pitch, yaw = att
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(roll), -np.sin(roll)],
                   [0, np.sin(roll), np.cos(roll)]])
    
    Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                   [0, 1, 0],
                   [-np.sin(pitch), 0, np.cos(pitch)]])
    
    Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                   [np.sin(yaw), np.cos(yaw), 0],
                   [0, 0, 1]])
    
    R = Rz @ Ry @ Rx
    
    # 应用旋转
    axes = axes @ R.T
    
    # 添加位置偏移
    axes += pos
    
    return axes

def create_quadcopter_faces(vertices):
    """创建四旋翼的面"""
    # 创建中心点
    center = np.mean(vertices, axis=0)
    
    # 创建连接中心点和各个顶点的三角形
    faces = []
    for i in range(4):
        faces.append([center, vertices[i], vertices[(i+1)%4]])
    
    return faces

# PID 控制器
class PID:
    def __init__(self, kp, ki, kd, setpoint=0):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.setpoint = setpoint
        self.integral = 0
        self.last_error = 0

    def update(self, value, dt):
        error = self.setpoint - value
        self.integral += error * dt
        derivative = (error - self.last_error) / dt
        self.last_error = error
        return self.kp * error + self.ki * self.integral + self.kd * derivative

    def set_setpoint(self, setpoint):
        self.setpoint = setpoint

# 状态变量
pos = np.array([0.0, 0.0, 0.0])   # 位置 x, y, z
vel = np.array([0.0, 0.0, 0.0])   # 速度 vx, vy, vz
att = np.array([0.0, 0.0, 0.0])   # 姿态角 roll, pitch, yaw
ang_vel = np.array([0.0, 0.0, 0.0])  # 角速度

# 飞行模式
current_mode = FlightMode.TAKEOFF
current_waypoint_index = 0  # 提前声明

# PID 控制器
pid_x = PID(2.0, 0.0, 1.0)
pid_y = PID(2.0, 0.0, 1.0)
pid_z = PID(10.0, 2.0, 5.0)

# 姿态控制器
pid_roll = PID(2.0, 0.0, 0.5)
pid_pitch = PID(2.0, 0.0, 0.5)
pid_yaw = PID(2.0, 0.0, 0.5)

# 确保起飞测试时有目标高度
if test_mode == TestMode.TAKEOFF_TEST:
    pid_z.set_setpoint(TAKEOFF_HEIGHT)
# 确保航路点测试时有目标航点
if test_mode == TestMode.WAYPOINT_TEST:
    current_mode = FlightMode.WAYPOINT
    current_waypoint_index = 0
    pid_x.set_setpoint(WAYPOINTS[0][0])
    pid_y.set_setpoint(WAYPOINTS[0][1])
    pid_z.set_setpoint(WAYPOINTS[0][2])

# 存储历史数据
positions = []
attitudes = []
flight_modes = []

def update_test_mode():
    global test_mode, current_mode, test_completion_time, pos, vel, att, ang_vel, current_waypoint_index
    
    if test_completion_time > 0:
        test_completion_time -= DT
        return
    
    if test_mode == TestMode.TAKEOFF_TEST:
        if current_mode == FlightMode.TAKEOFF and pos[2] >= TAKEOFF_HEIGHT:
            test_mode = TestMode.LANDING_TEST
            current_mode = FlightMode.LANDING
            pid_z.set_setpoint(0)
            test_completion_time = TEST_COMPLETION_TIME
    
    elif test_mode == TestMode.LANDING_TEST:
        if current_mode == FlightMode.LANDING and pos[2] <= 0.1:
            test_mode = TestMode.WAYPOINT_TEST
            current_mode = FlightMode.TAKEOFF
            current_waypoint_index = 0
            pid_x.set_setpoint(WAYPOINTS[0][0])
            pid_y.set_setpoint(WAYPOINTS[0][1])
            pid_z.set_setpoint(WAYPOINTS[0][2])
            test_completion_time = TEST_COMPLETION_TIME
    
    elif test_mode == TestMode.WAYPOINT_TEST:
        if current_mode == FlightMode.WAYPOINT and current_waypoint_index >= len(WAYPOINTS):
            test_mode = TestMode.RETURN_HOME_TEST
            current_mode = FlightMode.RETURN_HOME
            pid_x.set_setpoint(HOME_POSITION[0])
            pid_y.set_setpoint(HOME_POSITION[1])
            pid_z.set_setpoint(HOME_POSITION[2])
            test_completion_time = TEST_COMPLETION_TIME
    
    elif test_mode == TestMode.RETURN_HOME_TEST:
        if current_mode == FlightMode.RETURN_HOME and np.linalg.norm(pos - HOME_POSITION) < WAYPOINT_RADIUS:
            test_mode = TestMode.COMPLETE_TEST
            current_mode = FlightMode.LANDING
            pid_z.set_setpoint(0)
            test_completion_time = TEST_COMPLETION_TIME

def update_flight_mode():
    global current_mode, current_waypoint_index
    
    if current_mode == FlightMode.TAKEOFF:
        if pos[2] >= TAKEOFF_HEIGHT:
            if test_mode == TestMode.TAKEOFF_TEST:
                return
            current_mode = FlightMode.WAYPOINT
            current_waypoint_index = 0
            pid_x.set_setpoint(WAYPOINTS[0][0])
            pid_y.set_setpoint(WAYPOINTS[0][1])
            pid_z.set_setpoint(WAYPOINTS[0][2])
    
    elif current_mode == FlightMode.WAYPOINT:
        current_waypoint = WAYPOINTS[current_waypoint_index]
        distance = np.linalg.norm(pos - current_waypoint)
        
        if distance < WAYPOINT_RADIUS:
            current_waypoint_index += 1
            if current_waypoint_index >= len(WAYPOINTS):
                if test_mode == TestMode.WAYPOINT_TEST:
                    return
                current_mode = FlightMode.RETURN_HOME
                pid_x.set_setpoint(HOME_POSITION[0])
                pid_y.set_setpoint(HOME_POSITION[1])
                pid_z.set_setpoint(HOME_POSITION[2])
            else:
                next_waypoint = WAYPOINTS[current_waypoint_index]
                pid_x.set_setpoint(next_waypoint[0])
                pid_y.set_setpoint(next_waypoint[1])
                pid_z.set_setpoint(next_waypoint[2])
    
    elif current_mode == FlightMode.RETURN_HOME:
        distance = np.linalg.norm(pos - HOME_POSITION)
        if distance < WAYPOINT_RADIUS:
            if test_mode == TestMode.RETURN_HOME_TEST:
                return
            current_mode = FlightMode.LANDING
            pid_z.set_setpoint(0)
    
    elif current_mode == FlightMode.LANDING:
        # 让 setpoint 以最大下降速度递减
        new_setpoint = max(pid_z.setpoint - LANDING_SPEED * DT, 0)
        pid_z.set_setpoint(new_setpoint)
        if pos[2] <= 0.1:
            if test_mode == TestMode.COMPLETE_TEST:
                return
            current_mode = FlightMode.TAKEOFF

# 创建图形
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# 设置坐标轴范围
ax.set_xlim([-3, 3])
ax.set_ylim([-3, 3])
ax.set_zlim([-3, 3])
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

# 绘制航路点
for waypoint in WAYPOINTS:
    ax.scatter(waypoint[0], waypoint[1], waypoint[2], color='green', marker='^')

# 绘制返航点
ax.scatter(HOME_POSITION[0], HOME_POSITION[1], HOME_POSITION[2], color='red', label='Home')

# 创建四旋翼多边形
vertices = create_quadcopter_vertices(pos, att)
faces = create_quadcopter_faces(vertices)
quadcopter = Poly3DCollection(faces, alpha=0.5, facecolor='blue', edgecolor='black')
ax.add_collection3d(quadcopter)

# 创建姿态指示器
axes = create_attitude_indicators(pos, att)
x_axis, = ax.plot([pos[0], axes[0,0]], [pos[1], axes[0,1]], [pos[2], axes[0,2]], 'r-', linewidth=2)
y_axis, = ax.plot([pos[0], axes[1,0]], [pos[1], axes[1,1]], [pos[2], axes[1,2]], 'g-', linewidth=2)
z_axis, = ax.plot([pos[0], axes[2,0]], [pos[1], axes[2,1]], [pos[2], axes[2,2]], 'b-', linewidth=2)

# 创建轨迹线
trajectory_line, = ax.plot([], [], [], 'b-', alpha=0.5)

# 添加模式文本
mode_text = ax.text2D(0.02, 0.95, '', transform=ax.transAxes)

def update(frame):
    global pos, vel, att, ang_vel
    
    # 更新测试模式
    update_test_mode()
    
    # 更新飞行模式
    update_flight_mode()
    
    # 计算每个方向的推力
    fx = pid_x.update(pos[0], DT)
    fy = pid_y.update(pos[1], DT)
    fz = pid_z.update(pos[2], DT) + MASS * GRAVITY

    # 计算姿态控制力矩
    tau_roll = pid_roll.update(att[0], DT)
    tau_pitch = pid_pitch.update(att[1], DT)
    tau_yaw = pid_yaw.update(att[2], DT)

    # 合力
    force = np.array([fx, fy, fz])
    torque = np.array([tau_roll, tau_pitch, tau_yaw])

    # 加速度
    acc = force / MASS - np.array([0, 0, GRAVITY])
    ang_acc = torque / INERTIA

    # 更新速度和位置
    vel += acc * DT
    pos += vel * DT
    # 防止降到地面以下
    if pos[2] < 0:
        pos[2] = 0
        vel[2] = 0

    # 更新角速度和姿态
    ang_vel += ang_acc * DT
    att += ang_vel * DT

    positions.append(pos.copy())
    attitudes.append(att.copy())
    flight_modes.append(current_mode)

    # 更新四旋翼位置和姿态
    vertices = create_quadcopter_vertices(pos, att)
    faces = create_quadcopter_faces(vertices)
    quadcopter.set_verts(faces)

    # 更新姿态指示器
    axes = create_attitude_indicators(pos, att)
    x_axis.set_data([pos[0], axes[0,0]], [pos[1], axes[0,1]])
    x_axis.set_3d_properties([pos[2], axes[0,2]])
    y_axis.set_data([pos[0], axes[1,0]], [pos[1], axes[1,1]])
    y_axis.set_3d_properties([pos[2], axes[1,2]])
    z_axis.set_data([pos[0], axes[2,0]], [pos[1], axes[2,1]])
    z_axis.set_3d_properties([pos[2], axes[2,2]])

    # 更新轨迹线
    positions_array = np.array(positions)
    trajectory_line.set_data(positions_array[:, 0], positions_array[:, 1])
    trajectory_line.set_3d_properties(positions_array[:, 2])

    # 更新模式文本
    mode_text.set_text(f'Test Mode: {test_mode.name}\nFlight Mode: {current_mode.name}')

    return quadcopter, x_axis, y_axis, z_axis, trajectory_line, mode_text

# 创建动画
ani = FuncAnimation(fig, update, frames=int(SIM_TIME/DT), interval=DT*1000, blit=True)
plt.legend()
plt.show()